#!/usr/bin/env python
# -*- coding:utf-8 -*-

""" 
:Description: pyspark udf example
:Owner: leo_jie
:Create time: 2019/12/5
"""
import random
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import IntegerType, StringType, ArrayType

spark = SparkSession.builder.master("local") \
    .appName("leo_pyspark_learning") \
    .getOrCreate()

df = spark.createDataFrame([
    (1, 144.5, 5.9, 33, 'M'),
    (2, 167.2, 5.4, 45, 'M'),
    (3, 124.1, 5.2, 23, 'F'),
    (4, 144.5, 5.9, 33, 'M'),
    (5, 133.2, 5.7, 54, 'F'),
    (3, 124.1, 5.2, 31, 'F'),
    (5, 129.2, 5.3, 42, 'M'),
], ['id', 'weight', 'height', 'age', 'gender'])

"""
udf的源码如下，需要两个参数
f 为一个函数，是一个必须输入的参数
returnType是返回值类型，默认为StringType()类型
def udf(f=None, returnType=StringType()):
    ......
"""

# 最简单的udf，举例：对某一列的值加1

add_one_udf = udf(lambda c: c + 1)

# 使用的时候也很简单有多种方式可供选择

df.select("id", "age", add_one_udf(df.age).alias("age+1"))
df.select("id", "age", add_one_udf('age').alias("age+1"))
df.select("id", "age", add_one_udf(col("age")).alias("age+1"))

# 最终的输出结果

"""
+---+---+-----+
| id|age|age+1|
+---+---+-----+
|  1| 33|   34|
|  2| 45|   46|
|  3| 23|   24|
|  4| 33|   34|
|  5| 54|   55|
|  3| 31|   32|
|  5| 42|   43|
+---+---+-----+
"""


# 当然你还可以这样去定义udf函数，显式声明一个外部函数

def add_one(c):
    return c + 1


# 注册udf函数时指定返回值类型
add_one_udf2 = udf(add_one, IntegerType())

df.select("id", "age", add_one_udf2(df.age).alias("age+1"))

"""
+---+---+-----+
| id|age|age+1|
+---+---+-----+
|  1| 33|   34|
|  2| 45|   46|
|  3| 23|   24|
|  4| 33|   34|
|  5| 54|   55|
|  3| 31|   32|
|  5| 42|   43|
+---+---+-----+
"""


# 当然 下面还有更简洁的方式来声明一个udf，即装饰器

# @udf(returnType=StringType()) 当然这里也允许你指定返回值类型
@udf
def to_lower_udf(c):
    return c.lower()


df.select("id", "gender", to_lower_udf(df.gender).alias("lower_gender"))

# 最终的结果

"""
+---+------+------------+
| id|gender|lower_gender|
+---+------+------------+
|  1|     M|           m|
|  2|     M|           m|
|  3|     F|           f|
|  4|     M|           m|
|  5|     F|           f|
|  3|     F|           f|
|  5|     M|           m|
+---+------+------------+
"""


# 如果你需要更复杂的传参，那么，请往下看

def gender_to_chinese(gender, gender_dict):
    return gender_dict[gender]


def gender_to_chinese_udf(gender_dict):
    return udf(lambda c: gender_to_chinese(c, gender_dict), StringType())


df.select("id", "gender", gender_to_chinese_udf({'F': u'男', 'M': u'女'})(col("gender")).alias("chinese_gender"))
df.select("id", "gender", gender_to_chinese_udf({'F': u'男', 'M': u'女'})(df.gender).alias("chinese_gender"))
df.select("id", "gender", gender_to_chinese_udf({'F': u'男', 'M': u'女'})(df['gender']).alias("chinese_gender"))
df.select("id", "gender", gender_to_chinese_udf({'F': u'男', 'M': u'女'})("gender").alias("chinese_gender"))
# 最终结果

"""
这个udf的作用是，接受一个字典传参，把英文性别转换为中文性别
+---+------+--------------+
| id|gender|chinese_gender|
+---+------+--------------+
|  1|     M|            女|
|  2|     M|            女|
|  3|     F|            男|
|  4|     M|            女|
|  5|     F|            男|
|  3|     F|            男|
|  5|     M|            女|
+---+------+--------------+
"""


# 当然，你还以指定多个参数，甚至传多个复杂类型的参数

def people_get_salary(gender, gender_dict, salary_list):
    return [gender_dict[gender], str(random.sample(salary_list, 1)[0])]


def people_get_salary_udf(gender_dict, salary_list):
    return udf(lambda c: people_get_salary(c, gender_dict, salary_list), ArrayType(StringType()))


# 当然下面的声明方法更简洁

def people_get_salary_udf2(gender_dict, salary_list):
    def people_get_salary2(gender):
        return [gender_dict[gender], str(random.sample(salary_list, 1)[0])]

    return udf(people_get_salary2, returnType=ArrayType(StringType()))


# 这个udf的需求很简单，接受gender_dict参数把gender列英文性别转中文，接受salary_list参数，从中随机取一个元素作为这个人的工资，
# 最后返回一个数组


df.select("id", "gender", people_get_salary_udf({'F': u'男', 'M': u'女'}, [100, 200, 300])('gender').alias("salary_list"))

df.select("id", "gender",
          people_get_salary_udf2({'F': u'男', 'M': u'女'}, [100, 200, 300])('gender').alias("salary_list"))

# 最终结果
"""
+---+------+-----------+
| id|gender|salary_list|
+---+------+-----------+
|  1|     M|  [女, 200]|
|  2|     M|  [女, 100]|
|  3|     F|  [男, 100]|
|  4|     M|  [女, 300]|
|  5|     F|  [男, 300]|
|  3|     F|  [男, 100]|
|  5|     M|  [女, 200]|
+---+------+-----------+
"""

# 那么 问题来了，如果我们需要在selectExpr或spark.sql(select udf(field) from table)中直接使用自定义函数，是否可以呢？接着看
# 用我们上面注册的最简单的udf为例

# df.selectExpr("id", "age", "add_one_udf(age) as age_add").show()

# 异常如下
# pyspark.sql.utils.AnalysisException: u"Undefined function: 'add_one_udf'.

# 如果你还不死心，那么演示一个spark.sql查询临时表

df.createOrReplaceTempView("people_table")

# spark.sql("select id, age, add_one_udf(age) as age_add from people_table").show()

# 依然是这个异常

# pyspark.sql.utils.AnalysisException: u"Undefined function: 'add_one_udf'.

# 此时， 就需要我们的spark.udf.register出场了

spark.udf.register("add_one_udf", lambda c: c + 1, IntegerType())
# 同时也可以把创建好的udf函数注册
spark.udf.register("add_one_udf2", add_one_udf2)

df.selectExpr("id", "age", "add_one_udf(age) as age_add")
spark.sql("select id, age, add_one_udf2(age) as age_add from people_table")

# 最终的结果

"""
+---+---+-------+
| id|age|age_add|
+---+---+-------+
|  1| 33|     34|
|  2| 45|     46|
|  3| 23|     24|
|  4| 33|     34|
|  5| 54|     55|
|  3| 31|     32|
|  5| 42|     43|
+---+---+-------+
"""

"""
  简单看一下register的源码，需要传一个udf名，一个函数，还有返回值类型
  def register(self, name, f, returnType=None):
      ......
"""

# 如果是多参数，较为复杂的udf类型怎么注册呢？

gender_dic = {'F': u'男', 'M': u'女'}
salary_list = [100, 200, 300]


def people_get_salary2(gender, gender_dict_in=gender_dic, salary_list_in=salary_list):
    return [gender_dict_in[gender], str(random.sample(salary_list_in, 1)[0])]


spark.udf.register("people_get_salary_udf2", people_get_salary2, ArrayType(StringType(), True))

spark.sql("select id, age, people_get_salary_udf2(gender) as salary_list from people_table").show()

"""
+---+---+-----------+
| id|age|salary_list|
+---+---+-----------+
|  1| 33|  [女, 100]|
|  2| 45|  [女, 300]|
|  3| 23|  [男, 200]|
|  4| 33|  [女, 100]|
|  5| 54|  [男, 100]|
|  3| 31|  [男, 100]|
|  5| 42|  [女, 200]|
+---+---+-----------+
"""

spark.stop()
